package mrsimulator;

/**
 * JobTrackerSimulator simulates a real JobTracker of MapReduce system
 */

import log.PVLogger;
import pfsimulator.PNode;
import pfsimulator.ParaflowSimulator;
import pfsimulator.VNode;
import util.IdFactory;

import java.util.*;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingDeque;

public class JobTrackerSimulator implements HeartbeatMonitor {
    public static PVLogger LOG = PVLogger.getLogger(JobTrackerSimulator.class);
    private static Random dice = new Random(System.currentTimeMillis());
    // Number of physical machine in the cluster
    private static final int CLUSTER_SIZE = 2;
    // Number of virtual machine on each host
    private static final int VM_PER_HOST = 2;
    // NIC bandwidth, in byte
    private static final long NIC_BW = 1024*1024*1024; // 1Gbyte
    // Number of NIC on each physical host
    private static final int NIC_NUM = 4;

    private PNode[] pNodes = null;
    private int numOfPNode = 0;
    private int numOfVNode = 0;
    private int numOfVNodePerPNode = 0;
    private final JobInProgressQM jipQm = new JobInProgressQM();
    private JobSubmissionQM submissionQM = null;
    private Scheduler scheduler;
    private ParaflowSimulator transferManager = new ParaflowSimulator();

    // Mappings
    private Map<Long, PNode> pNodeMap = null; // PNode id --> PNode
    private Map<String, TaskTrackerSimulator> taskTrackerSimulatorMap = null; // VNodeId --> TT
   // private Map<String, TaskInProgress> runningMapIdToTIP; // TaskId --> TaskInProgress
    private Map<String, TaskInProgress> runningReduceIdToTIP; // TaskId --> TaskInProgress

    /**
     * Constructor
     */
    public JobTrackerSimulator() {
        this.submissionQM = new JobSubmissionQM();
        this.runningReduceIdToTIP = new HashMap<String, TaskInProgress>();
    //    this.runningMapIdToTIP = new HashMap<String, TaskInProgress>();
        this.taskTrackerSimulatorMap = new HashMap<String, TaskTrackerSimulator>();
        this.pNodeMap = new HashMap<Long, PNode>();
        this.createCluster(CLUSTER_SIZE, VM_PER_HOST, NIC_BW * NIC_NUM);
        // Set scheduler
        // TODO:FIXME --> replace DumbScheduler with a real scheduler
        this.scheduler = new DumbScheduler(this);
    }
    public void offerService(){
        Collection<TaskTrackerSimulator> allTT = taskTrackerSimulatorMap.values();
        // start all task tracker
        for(TaskTrackerSimulator tt : allTT){
            tt.start();
        }
        this.submissionQM.start();
    }

    //
    // Internal functions
    //
    private void addPNode(long id, PNode pNode){
        pNodeMap.put(id, pNode);
    }

    private void createTaskTrackerForVNode(VNode vNode){
        taskTrackerSimulatorMap.put(vNode.getHostName(), new TaskTrackerSimulator(vNode.getHostName(), this));
    }

    // Clear up when job is done
    private void jobClearUp(String jobId) {
        JobInProgress jobDone = jipQm.getJobInProgress(jobId);
        LOG.info("Job " + jobId + " done. ExecutionTime = " +
             (System.currentTimeMillis() - jobDone.getStartTime()) + " seconds");
        jipQm.deleteJIP(jobId);
    }

    //
    // APIs
    //

    /*
     * Workload generator calls this function to submit job
     */
    public synchronized void submitJob(Job job){
        LOG.debug("Job " + job.getJobId() + " is submitted");
        this.submissionQM.addJob(job);
    }

    /**
     * Move head of JobInProgress Queue to the tail
     * Thread NOT safe
     * */
    public void roundRobin(){
        JobInProgress head = jipQm.removeQHead();
        if(head != null){
            jipQm.pushToTail(head);
        }
    }

    public List<JobInProgress> getJobQueue(){
        return this.jipQm.getJobQueue();
    }


    public TaskTrackerAction heartbeat(TaskTrackerStatus status){
        // Lock on JobInProgress queue to make it thread safe
        TaskTrackerAction action = new TaskTrackerAction();
        synchronized (jipQm) {
            // First process finished tasks and update internal status, then schedule new tasks
            // - Get map task done events, mark corresponding maps as done,
            //   -- if all maps are done, mark notification on reduce
            // - Get reduce task done event, mark reduce task as done,check
            //  -- if all reduces of a job are done, report a job finished event
            // - Call scheduler to assign task
            // - Check all running reduce on this tt, prepare for notification if necessary

            // process finished map tasks
            enqueueFlowTask(status.getMapFinishReport());

            // process finished reduce tasks
            List<Task> finishedReduceTasks = status.getReduceFinishReport();
            boolean jobDone = false;
            for(Task reduce: finishedReduceTasks){
                jobDone = jipQm.processReduceDoneEvent(reduce);
                //remove from bookkeeping
                this.runningReduceIdToTIP.remove(reduce.getTaskId());
                if(jobDone){
                    jobClearUp(reduce.getJobId());
                }
            }

            // Check all running reduce on this tt and prepare for notification if necessary
            List<Task> runningReduceTask = status.getRunningReduceTasks();
            List<Task> reduceToNotify = new LinkedList<Task>();
            for(Task t: runningReduceTask){
                TaskInProgress runningReduce = this.runningReduceIdToTIP.get(t.getTaskId());
                if(runningReduce != null){
                    if(runningReduce.shouldNotifyReduce()){
                        reduceToNotify.add(t);
                        runningReduce.markReduceAsNotified();
                    }
                } else {
                    LOG.error("Cannot find running reduce task " + t.getTaskId());
                }
            }

            // Call scheduler to assign task
            // If reduce is launched, add it to bookkeeping
            List<TaskInProgress> tasksToLaunch = scheduler.assignTask(status);
            List<Task> mapTasksToLaunch = new LinkedList<Task>();
            List<Task> reduceTasksToLaunch = new LinkedList<Task>();
            for(TaskInProgress tip: tasksToLaunch){
                // separate maps and reduces
                if(tip.isMap()){
                    mapTasksToLaunch.add(tip.getTask());
                    //this.runningMapIdToTIP.put(tip.getTaskId(), tip);
                } else {
                    reduceTasksToLaunch.add(tip.getTask());
                    this.runningReduceIdToTIP.put(tip.getTaskId(), tip);
                }
            }

            // Fill in the heartbeat response
            //  maps, reduces, reduce to notify
            action.addMapTasks(mapTasksToLaunch);
            action.addReduceTasks(reduceTasksToLaunch);
            action.addReduceTasksToNotify(reduceToNotify);
        }
        return action;
    }

    //TODO: related to flow/////////////////////////--->
    private void enqueueFlowTask(List<Task> finishedMaps) {
        // TODO
        // - For each finished map task in the list, enqueue a flow task
        // - specify
    }

    private class FlowTask{
        private String mapId;
        private String reduceId;
        private String flowId;
    }

    //TODO//////////////////////////////////////////<---

    /**
     * Create a cluster with given num of PNodes, each of witch
     * has numOfVNodePerPNode on it and bwPerPNode bytes of bandwidth
     * */
    public void createCluster(int numOfPNode, int numOfVNodePerPNode, long bwPerPNode){
        this.numOfPNode = numOfPNode;
        this.numOfVNode = numOfPNode*numOfVNodePerPNode;
        this.numOfVNodePerPNode = numOfVNodePerPNode;
        for (int i = 0; i < numOfPNode; i++) {
            long pNodeId = IdFactory.getPNodeId();
            PNode pNode = new PNode(pNodeId, bwPerPNode);
            for (int j = 0; j < numOfPNode; j++) {
                String vid = IdFactory.getVNodeId(pNodeId);
                VNode vNode = new VNode(vid, IdFactory.getNextIp(), pNode);
                createTaskTrackerForVNode(vNode);
                pNode.addVNode(vNode);
            }
            addPNode(pNodeId, pNode);
        }
        // Let's assume that the PNode set never changes, we keep
        // PNodes in array for better use
        this.pNodes = pNodeMap.values().toArray(new PNode[0]);
        // Initialize the pfsimulator
        this.transferManager.construct(new LinkedList<PNode>(pNodeMap.values()));
    }

    // Randomly pick a VNode from the cluster
    private VNode getRandomNode(){
        PNode pNode = pNodes[dice.nextInt(pNodes.length)];
        return pNode.getRandomVNode();
    }

    //
    // Internal classes
    //

    // Job submission QM keeps a queue of submitted job, and create JonInProgress
    // when new Job is added to queue
    private class JobSubmissionQM {
        private BlockingQueue<Job> jobSubmissions = new LinkedBlockingDeque<Job>();
        private Thread initializer = new Thread(new JobInitializer());

        public void addJob(Job newJob){
            try{
                LOG.debug("Job " + newJob.getJobId() + "added to submission queue");
                jobSubmissions.put(newJob);
            } catch (InterruptedException e){
                LOG.error("Swallow interruption exception when submitting job");
            }
        }

        class JobInitializer implements Runnable {

            @Override
            public void run() {
                try{
                    while(!Thread.interrupted()){
                        initJob(jobSubmissions.take());
                    }
                } catch (InterruptedException e){
                    LOG.warn("Job initializer interrupted");
                }
            }

            private void initJob(Job newJob){
                LOG.debug("Initializing job " + newJob.getJobId());
                JobInProgress jip = new JobInProgress(newJob);
                synchronized (jipQm){
                    jipQm.addJIP(jip);
                }
            }
        }

        // Call this function to start listener thread
        public void start(){
            this.initializer.start();
        }
    }

    /**
     * A class that groups all the information from a JIP that is necessary for
     * scheduling a job
     * */
    static class JobSchedulingContext {

    }

    // JobInProgressQM keeps states of all running jobs. It works as follows:
    // - Accept new JIP from JobSubmissionQM
    // - Keep basic FIFO order
    // - Keep update JIP status when task finished, delete if needed
    private class JobInProgressQM {
        private List<JobInProgress> jobQueue = new LinkedList<JobInProgress>();

        public void addJIP(JobInProgress jip){
            jobQueue.add(jip);
        }
        public void deleteJIP(String jobId){
            for (int i = 0; i < jobQueue.size(); i++) {
                JobInProgress jobInProgress = jobQueue.get(i);
                if(jobId.equals(jobInProgress.getJobId())){
                    jobQueue.remove(i);
                }
            }
        }

        public JobInProgress removeQHead(){
            if(jobQueue.isEmpty()) return null;
            return jobQueue.remove(0);
        }

        public void pushToTail(JobInProgress jip){
            jobQueue.add(jobQueue.size(), jip);
        }

        /**
         * Thread NOT safe.
         * Return a copy of the job queue
         * */
        public List<JobInProgress> getJobQueue(){
            return jobQueue;
        }

        /*
         * Update the internal bookkeeping when shuffle of map is done, which
         * means the end of a map task.
         * If all map tasks ends, mark reduce tasks to be notified
         *
         * Return true if this is the last running map of the job it belongs to;
         * else return false
         */
        public boolean processShuffleDoneEvent(Task map){
            String jobId = map.getJobId();
            String taskId = map.getTaskId();
            if(!map.isMap()){
                LOG.error("Calling the wrong update map function " + taskId);
                return false;
            }
            boolean allMapTaskDone = false;
            JobInProgress jobInProgress = getJobInProgress(jobId);
            if(jobInProgress == null){
                LOG.error("No such job " + jobId);
            } else {
                allMapTaskDone = jobInProgress.updateRunningMap(map);
                if(allMapTaskDone){
                    jobInProgress.notifyAllReduceTaskToRun();
                }
            }
            return allMapTaskDone;
        }

        // TODO: mark task as finished and
        public void processMapDoneEvent(Task map){
        }

        /**
         *
         * */

        public boolean processReduceDoneEvent(Task reduce){
            String jobId = reduce.getJobId();
            String taskId = reduce.getTaskId();
            if(reduce.isMap()){
                LOG.error("Calling the wrong update reduce function " + taskId);
                return false;
            }

            boolean allReduceDone = false;
            JobInProgress jobInProgress = getJobInProgress(jobId);
            if(jobInProgress == null){
                LOG.error("No such job " + jobId);
            } else {
                allReduceDone = jobInProgress.updateRunningReduce(reduce);
            }
            /*
            for (int i = 0; i < jobQueue.size(); i++) {
                JobInProgress jobInProgress = jobQueue.get(i);
                if(jobInProgress.getJobId().equals(jobId)){
                    allReduceDone = jobInProgress.updateRunningReduce(reduce);
                    break;
                }
            }
            */
            return allReduceDone;
        }

        /**
         * Thread *NOT* safe, synchronize JIPQM before calling
         * */
        public JobInProgress getJobInProgress(String jobId){
            for (int i = 0; i < jobQueue.size(); i++) {
                JobInProgress jobInProgress = jobQueue.get(i);
                if(jobInProgress.getJobId().equals(jobId))
                    return jobInProgress;
            }
            return null;
        }

        public void processMapShuffle(Task pendingMap){
        }
    }

    public class JobInProgress {

        private String jobId;
        private TaskInProgress[] maps = null;
        private TaskInProgress[] reduces = null;
        private int nonRunningMapTaskNum;
        private int nonRunningReduceTaskNum;
        private int runningMapTaskNum;
        private int runningReduceTaskNum;
        private int finishedMapTaskNum;
        private int finishedReduceTaskNum;

        private long startTime;

        public JobInProgress(Job job) {
            this.jobId = job.getJobId();
            this.maps = new TaskInProgress[job.getNumOfMaps()];
            this.reduces = new TaskInProgress[job.getNumOfReduces()];
            createTasks(job);
            nonRunningMapTaskNum = maps.length;
            nonRunningReduceTaskNum = reduces.length;
            runningMapTaskNum = 0;
            runningReduceTaskNum = 0;
            finishedMapTaskNum = 0;
            finishedReduceTaskNum = 0;
            startTime = System.currentTimeMillis();
        }

        public String getJobId() {
            return jobId;
        }

        private void createTasks(Job job) {
            // create map tip
            Task[] allMaps = job.getAllMaps();
            for (int i = 0; i < allMaps.length; i++) {
                maps[i] = new TaskInProgress(allMaps[i]);
            }
            // create reduce tip
            Task[] allReduces = job.getAllReduces();
            for (int i = 0; i < allReduces.length; i++) {
                reduces[i] = new TaskInProgress(allReduces[i]);
            }
        }

        public int getNonRunningMapTaskNum() {
            return nonRunningMapTaskNum;
        }

        public int getNonRunningReduceTaskNum() {
            return nonRunningReduceTaskNum;
        }

        public int getRunningMapTaskNum() {
            return runningMapTaskNum;
        }

        public int getRunningReduceTaskNum() {
            return runningReduceTaskNum;
        }

        public int getFinishedMapTaskNum() {
            return finishedMapTaskNum;
        }

        public int getFinishedReduceTaskNum() {
            return finishedReduceTaskNum;
        }

        public long getStartTime() {
            return startTime;
        }

        public TaskInProgress getTaskInProgress(Task task){
            TaskInProgress[] candidates = task.isMap() ? maps : reduces;
            String taskId = task.getTaskId();
            for(int i = 0; i < candidates.length; i++) {
                if(candidates[i].getTaskId().equals(taskId)){
                    return candidates[i];
                }
            }
            return null;
        }

        /*
         * Mark given task as done, return true if it's the last running map task\
         * else, return false
         */
        public boolean updateRunningMap(Task map){

            TaskInProgress taskInProgress = getTaskInProgress(map);
            // Sanity checking
            if(taskInProgress == null){
                LOG.error("No such task " + map.getTaskId());
                return false;
            }
            if(!taskInProgress.isTaskLaunched()){
                LOG.error("Task " + map.getTaskId() +" is not launched");
                return false;
            }
            if(taskInProgress.isTaskDone()){
                LOG.error("Task " + map.getTaskId() +" has already be marked as finished");
                return false;
            }
            // Status is running if u r here

            taskInProgress.markTaskAsDone();
            LOG.info("Job " + jobId + "'s map task " + map.getTaskId() + " is done");
            this.runningMapTaskNum--;
            this.finishedMapTaskNum++;

            if(this.maps.length == this.finishedMapTaskNum){
                LOG.debug("All map tasks of " + map.getJobId() + " are done");
                return true;
            } else {
                LOG.debug("Finished task number of " + map.getJobId() + " is " + this.finishedMapTaskNum + " where total map num is " + maps.length);
                return false;
            }
        }

        /*
         * Mark given reduce as done, return true if it's the last running reduce task,
         * aka, the whole job is done; else, return false
         */
        public boolean updateRunningReduce(Task reduce){

            TaskInProgress taskInProgress = getTaskInProgress(reduce);
            // Sanity check
            if(taskInProgress == null){
                LOG.error("No such task " + reduce.getTaskId());
                return false;
            }
            if(taskInProgress.isTaskLaunched() == false){
                LOG.error("Task " + reduce.getTaskId() +" is not launched");
                return false;
            }
            if(taskInProgress.isTaskDone() == true){
                LOG.error("Task " + reduce.getTaskId() +" has already be marked as finished");
                return false;
            }
            // Status is running if u r here

            taskInProgress.markTaskAsDone();
            LOG.info("Job " + jobId + "'s reduce task " + reduce.getTaskId() + " is done");
            this.runningReduceTaskNum--;
            this.finishedReduceTaskNum++;

            if(this.reduces.length == this.finishedReduceTaskNum){
                return true;
            } else {
                return false;
            }
        }

        /**
         * Notify reduce tasks that all map tasks are done, should run the real reduce function
         * */
        public void notifyAllReduceTaskToRun(){
            for (int i = 0; i < reduces.length; i++) {
                reduces[i].markMapPhaseAsDone();
            }
        }

        /**
         * Choose a non-running map task from this job, return null if
         * all map task have been launched
         * */
        public TaskInProgress pickMapTaskToRun(TaskTrackerStatus status) {
            // Choose a local task
            for(int i = 0; i < maps.length; i++){
                if(maps[i].isTaskLaunched() == false &&
                        maps[i].getSimuInputLocation().getHostName().equals(status.getHost())){
                    return maps[i];
                }
            }

            // Choose a non-local task
            for(int i = 0; i < maps.length; i++){
                if(!maps[i].isTaskLaunched()){
                    return maps[i];
                }
            }

            // all map tasks have been launched, return null
            return null;
        }

        /**
         * Choose a non-running reduce taks from this job, return null if all
         * reduce tasks have been launched
         * */

        public TaskInProgress pickReduceTaskToRun(TaskTrackerStatus status) {
            for(int i = 0; i < reduces.length; i++){
                if(!reduces[i].isTaskLaunched()){
                    return reduces[i];
                }
            }
            return null;
        }

        /*
        public void markShufflableMap(Task pendingMap) {
            for(TaskInProgress reduce: reduces){
                reduce.setShuffleStates(pendingMap.getTaskId(), ShuffleState.READY);
            }
        }
        */
    }

    private enum ShuffleState {NOT_READY, READY, IN_PROGRESS, DONE, UNKNOWN}

    public class TaskInProgress {
        private String taskId = null;
        private boolean isMap = true;
        private long simuTime = 0;
        private long simuOutputSize = 0;
        private VNode simuInputLocation = null;
        private boolean taskLaunched;
        private boolean taskDone;
        private Task task;
        private long startTime;

        // for reduce task only
        private boolean shouldNotifyReduce;
        private boolean notified;
        //private Map<String, ShuffleState> shuffleStates = new HashMap<String, ShuffleState>();

        public TaskInProgress(Task task){
            this.task = task;
            this.taskId = task.getTaskId();
            this.isMap = task.getTaskInfo().isMap();
            this.simuTime = task.getTaskInfo().getCompletionTime();
            this.simuOutputSize = task.getTaskInfo().getOutputSize();
            this.simuInputLocation = getRandomNode();
            this.taskLaunched = false;
            this.taskDone = false;
            this.shouldNotifyReduce = false;
            this.notified = false;
        }

        /**
         * Call this function whenever creating reduce tasks
        public void initReduce(List<Task> maps){
            for (Task map : maps) {
                this.shuffleStates.put(map.getTaskId(), ShuffleState.NOT_READY);
            }
        }

        public ShuffleState getShuffleState(String mapTaskId){
            if(!this.isMap()){
                ShuffleState state = shuffleStates.get(mapTaskId);
                if(state != null){
                    return state;
                }
            }
            return ShuffleState.UNKNOWN;
        }

        public void setShuffleStates(String mapId, ShuffleState newState){
            shuffleStates.put(mapId, newState);
        }
         * */

        public Task getTask() {
            return task;
        }

        public boolean isTaskDone() {
            return taskDone;
        }

        public void markTaskAsDone() {
            this.taskDone = true;
        }

        public boolean isTaskLaunched() {
            return taskLaunched;
        }

        public void markTaskAsLaunched() {
            this.taskLaunched = true;
            this.startTime = System.currentTimeMillis();
        }

        public String getTaskId() {
            return taskId;
        }

        public VNode getSimuInputLocation() {
            return simuInputLocation;
        }

        public boolean isMap(){
            return isMap;
        }

        //
        // Functions about reduce notification -->
        //
        public boolean shouldNotifyReduce(){
            return shouldNotifyReduce && !notified;
        }

        public void markMapPhaseAsDone(){
            this.shouldNotifyReduce = true;
        }

        public boolean isThisReduceNotified(){
            return this.notified;
        }

        public void markReduceAsNotified(){
            this.notified = true;
        }
        //
        // <-- Functions about reduce notification
        //
    }
}
